package org.shx.springboot.ollama.ai.config;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import reactor.core.scheduler.Schedulers;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 非必须，使用ollamaChatClient才需要
 *     @Qualifier("ollamaChatClient")
 *     private ChatClient ollamaChatClient;
 *
 *  可以直接使用
 *      @Autowired
 *     private OllamaChatModel chatModel;
 */
@Configuration
public class ChatClientConfig {
    @Resource
    private OllamaChatModel ollamaChatModel;

    //不需要对话历史记忆功能
    @Bean(name = "ollamaChatClient2")
    public ChatClient ollamaChatClient2() {
        // 移除所有与memory相关的配置
        return ChatClient.builder(ollamaChatModel)
                .defaultSystem("你好.")
                .build();
    }

    @Bean(name = "ollamaChatClient")
    public ChatClient ollamaChatClient() {
        // 1. 完整实现ChatMemory接口的所有抽象方法
        ChatMemory chatMemory = new ChatMemory() {
            // 使用Map存储多个对话的消息，key为conversationId
            private final Map<String, List<Message>> conversationMessages = new HashMap<>();

            // 新增的必要实现：添加消息列表到指定对话
            @Override
            public void add(String conversationId, List<Message> messages) {
                conversationMessages.computeIfAbsent(conversationId, k -> new ArrayList<>())
                        .addAll(messages);
            }

            // 添加单条消息到指定对话
            @Override
            public void add(String conversationId, Message message) {
                conversationMessages.computeIfAbsent(conversationId, k -> new ArrayList<>())
                        .add(message);
            }

            // 获取指定对话的所有消息
            @Override
            public List<Message> get(String conversationId) {
                return new ArrayList<>(conversationMessages.getOrDefault(conversationId, new ArrayList<>()));
            }

            // 清除指定对话的消息
            @Override
            public void clear(String conversationId) {
                conversationMessages.remove(conversationId);
//                conversationMessages.clear();

            }
        };

        // 不在这里固定conversationId，留到使用时指定
        // 2. 使用正确的Builder方法配置（根据源码使用conversationId而非defaultConversationId）
        MessageChatMemoryAdvisor memoryAdvisor = MessageChatMemoryAdvisor.builder(chatMemory)
//                .conversationId("default-conversation-id") // 修正方法名
                .order(0)
                .scheduler(Schedulers.boundedElastic())
                .build();


        // 3. 构建ChatClient
        return ChatClient.builder(ollamaChatModel)
                .defaultSystem("使用中文回复.")
                .defaultAdvisors(memoryAdvisor)
                .build();
    }
}
